from data import Horse_Colic
from torch.utils import data
import torch.nn as nn
import torch
from model import Logistic
import config
import torchnet
import visdom
import eval

# config 参数
opt = config.DefaultConfig()

# 模型加载
seed = 19970124
model = Logistic(input_dim=21)
torch.manual_seed(seed)


# 数据集加载
train_dataset = Horse_Colic(opt.data_root, train=True)
test_dataset = Horse_Colic(opt.data_root, train=False)
data_len = len(train_dataset)
train_dataloader = data.DataLoader(
    train_dataset,
    batch_size=opt.batch_size,
    shuffle=opt.shuffle,
    num_workers=opt.num_workers
)
test_dataloader = data.DataLoader(
    test_dataset,
    batch_size=len(test_dataset),
    shuffle=False,
    num_workers=opt.num_workers
)
# 损失函数和优化器
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=opt.lr, weight_decay=opt.weight_decay)

# 统计指标
loss_meter = torchnet.meter.AverageValueMeter()
confusion_matrix = torchnet.meter.ConfusionMeter(2)

# 训练过程
if __name__ == "__main__":
    # visdom
    viz = visdom.Visdom(env="Logistic")
    name = ["loss", "mean_loss", "acc"]
    for epoch in range(opt.max_epoch):
        loss_meter.reset()                      # 清空loss累计
        confusion_matrix.reset()                # 重置混淆矩阵
        acc = 0
        best_acc = 0.83
        for i, (data, label) in enumerate(train_dataloader):
            out = model(data)  # 计算输出
            optimizer.zero_grad()  # 清空梯度
            loss = criterion(out, label)  # 计算损失
            loss.backward()  # 反向传播
            optimizer.step()  # 更新学习率
            if i % 3 == 0:
                acc = eval.Eval(model, test_dataloader)
                if acc > best_acc:
                    best_acc = acc
                    torch.save(model, opt.save_path+"model.pkl")
            loss_meter.add(loss.data)
            viz.line(Y=[[float(loss.data), float(loss_meter.mean), acc*10]],
                     X=[[epoch + i * opt.batch_size / data_len for j in range(3)]],
                     win='line',
                     opts=dict(legend=name,
                               title='line test',
                               width=800,
                               height=800,
                               xlabel='Time',
                               ylabel='Volume'),
                     update='append'
                     )


